'''
Company: TWL
Author: xue jian
Email: xuejian@kanzhun.com
Date: 2020-11-16 14:31:27
'''
#
# @lc app=leetcode.cn id=1424 lang=python3
#
# [1424] 对角线遍历 II
#

# @lc code=start
from typing import List
class Solution:
    def findDiagonalOrder(self, nums: List[List[int]]) -> List[int]:
        ans = []
        import collections
        store = collections.OrderedDict()
        for i in range(len(nums)):
            for j in range(len(nums[i])):
                if i+j not in store:
                    store[i+j]=[]
                store[i+j].append(nums[i][j])
        for _, v in store.items():
            ans.extend(v[::-1])
        return ans
# @lc code=end

